from openai import AsyncOpenAI 
from agents import OpenAIChatCompletionsModel, Agent, Runner, set_default_openai_client 
from agents.mcp  import MCPServer, MCPServerStdio
from agents.model_settings  import ModelSettings
import asyncio 
from contextlib import AsyncExitStack
from agents import set_tracing_disabled

OPENAI_API_KEY = "hk-xxx"
OPENAI_API_BASE = "https://api.openai-hk.com/v1" 
MODEL = "deepseek-v3"
 
external_client = AsyncOpenAI(
    base_url=OPENAI_API_BASE,
    api_key=OPENAI_API_KEY,
)
 
set_default_openai_client(external_client)
set_tracing_disabled(True)
deepseek_model = OpenAIChatCompletionsModel(
    model=MODEL,
    openai_client=external_client)


async def mcp_run_multi(servers_params, message):
    # 使用 AsyncExitStack 自动管理多个上下文退出
    async with AsyncExitStack() as stack:
        servers = []
        # 创建并进入所有 server 上下文 
        for p in servers_params:
            server = MCPServerStdio(
                name=p.get("name", "Unnamed Server"),
                cache_tools_list=True,
                params={
                    "command": "uv",
                    "args": ["run", p["script"]],
                },
            )
            entered_server = await stack.enter_async_context(server)
            servers.append(entered_server)

            # 构造 agent，传入多个 server
        agent = Agent(
            name="Assistant",
            instructions=(
                "你是一名助人为乐的助手。请先调用query_weather工具查询北京天气，"
                "然后将查询结果通过write_file工具写入res.md文件中。"
            ),
            mcp_servers=servers,
            model_settings=ModelSettings(tool_choice="required"),
            model=deepseek_model
        )

        print(f"Running: {message}")
        result = await Runner.run(starting_agent=agent, input=message)
        print(result.final_output)

        return result

if __name__ == "__main__":
    # 调用：传入多个 server 的配置 
    result = asyncio.run(mcp_run_multi( 
        servers_params=[
            {"name": "Weather Server", "script": "weather_server.py"}, 
            {"name": "Writer Server", "script": "write_server.py"} 
        ],
        message="请帮我查询北京和南京天气，并将查询的结果写入本地res.md文档。"
    ))
    print(result)